/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#include <vector>
#include <cstdint>
#include <cmath>
#include <random>
#include <iostream>
#include <algorithm>
#include "acl/acl.h"
#include "atvc.h"

#define CHECK_ACL(x)                                                                        \
    do {                                                                                    \
        aclError __ret = x;                                                                 \
        if (__ret != ACL_ERROR_NONE) {                                                      \
            std::cerr << __FILE__ << ":" << __LINE__ << " aclError:" << __ret << std::endl; \
        }                                                                                   \
    } while (0)


namespace {
static constexpr float REL_TOL = 1e-3f;
static constexpr float ABS_TOL = 1e-5f;


bool IsClose(float a, float b)
{
    const float eps = 1e-40f;
    float diff = std::abs(a - b);
    return (diff <= ABS_TOL) || (diff <= REL_TOL * std::max(std::abs(a), std::abs(b) + eps));
}

bool VerifyResults(const std::vector<float> &golden, const std::vector<float> &output)
{
    for (int32_t i = 0; i < golden.size(); i++) {
        if (!IsClose(golden[i], output[i])) {
            printf("[ERROR]: Accuracy verification failed! The expected value of element "
                   "in index [%d] is %f, but actual value is %f.\n",
                i,
                golden[i],
                output[i]);
            return false;
        }
    }
    return true;
}

void InitializeACL(aclrtContext &context, aclrtStream &stream, int32_t deviceId)
{
    CHECK_ACL(aclInit(nullptr));
    CHECK_ACL(aclrtSetDevice(deviceId));
    CHECK_ACL(aclrtCreateContext(&context, deviceId));
    CHECK_ACL(aclrtCreateStream(&stream));
}

void CleanACL(aclrtStream &stream, aclrtContext &context, int32_t deviceId)
{
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtDestroyContext(context));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
}


using OP_TRAITS = ATVC::OpTraits<ATVC::OpInputs<float, float>, ATVC::OpOutputs<float>, ATVC::OpTemps<float, float>>;

template<typename Traits>
struct LeakyReluGradComputeFunc {
    template<typename T>
    __aicore__ inline void operator()(AscendC::LocalTensor<T> dy, AscendC::LocalTensor<T> x,
        AscendC::LocalTensor<T> dx, AscendC::LocalTensor<T> mask, AscendC::LocalTensor<T> tmp, float alpha) {
        
        /*
        AscendC::CompareScalar(temp, x, 0.0f, AscendC::CMPMODE::GT, dx.GetSize());
        AscendC::DataCopy(tmp1, dy, dx.GetSize());
        AscendC::Muls(tmp1, tmp1, alpha, dx.GetSize());
        AscendC::Select(dx, temp, dy, tmp1, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, dx.GetSize());
        */
        int32_t tiled_cnt = dx.GetSize();
        // 1. dx = dy * α（负半轴底量）
        // 1. 计算掩码：mask = (x>0) ? 1 : 0
        AscendC::Sign(mask, x, tiled_cnt);     // mask = sign(x) → {-1, 0, 1}
        AscendC::Relu(mask, mask, tiled_cnt);  // mask = relu(mask) → {0, 1}，专用掩码变量
        // 2. 计算负半轴梯度：dx = dy * alpha
        AscendC::Muls(dx, dy, alpha, tiled_cnt);
        // 3. 计算正半轴补充量：tmp = dy * (1 - alpha) * mask
        AscendC::Muls(tmp, dy, 1.0f - alpha, tiled_cnt);  // tmp = dy * (1 - alpha)
        AscendC::Mul(tmp, tmp, mask, tiled_cnt);           // tmp = dy * (1 - alpha) * mask
        // 4. 合并结果：正半轴得到dy，负半轴保持dy*alpha
        AscendC::Add(dx, dx, tmp, tiled_cnt);
      
    }
};

void InitializeData(int32_t eleNum, std::vector<float> &dy, std::vector<float> &x, std::vector<float> &golden,
    float alpha)
{
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dis(-100.0f, 100.0f);  

    for (int i = 0; i < eleNum; ++i) {
            dy[i] = dis(gen);
            x[i] = dis(gen);
            golden[i] = (x[i] > 0) ? dy[i] : dy[i] * alpha;
    }
}

void CleanUp(uint8_t *&dxHost, uint8_t *&dyDevice, uint8_t *&xDevice, uint8_t *&dxDevice)
{
    CHECK_ACL(aclrtFree(dyDevice));
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(dxDevice));
    CHECK_ACL(aclrtFreeHost(dxHost));
}
}


template<class Traits>
__global__ __aicore__ void LeakyReluGradCustom(GM_ADDR dy, GM_ADDR x, GM_ADDR dx, ATVC::EleWiseParam param, float alpha)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    auto op = ATVC::Kernel::EleWiseOpTemplate<LeakyReluGradComputeFunc<Traits>>();
    op.Run(dy, x, dx, &param, alpha);
}

int main()
{
    int32_t eleNum = 13 * 1333;
    size_t inputByteSize = static_cast<size_t>(eleNum) * sizeof(float);
    size_t outputByteSize = static_cast<size_t>(eleNum) * sizeof(float);
    std::vector<float> dy(eleNum);
    std::vector<float> x(eleNum);
    std::vector<float> golden(eleNum);
    float alpha = 0.1f;  

    InitializeData(eleNum, dy, x, golden, alpha);

    aclrtContext context;
    aclrtStream stream = nullptr;
    int32_t deviceId = 0;
    InitializeACL(context, stream, deviceId);

    ATVC::EleWiseParam param;
    if (!ATVC::Host::CalcEleWiseTiling<OP_TRAITS>(eleNum, param)) {
            printf("Elewise tiling error.\n");
            return -1;
    };
    auto elementParamSize = sizeof(param);

    uint8_t *dxHost;
    uint8_t *dyDevice;
    uint8_t *xDevice;
    uint8_t *dxDevice;

    CHECK_ACL(aclrtMallocHost((void **)&dxHost, outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&dyDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&dxDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    CHECK_ACL(aclrtMemcpy(dyDevice, inputByteSize, dy.data(), inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, x.data(), inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    LeakyReluGradCustom<OP_TRAITS>
        <<<param.tilingData.blockNum, nullptr, stream>>>(dyDevice, xDevice, dxDevice, param, alpha);

    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtMemcpy(dxHost, outputByteSize, dxDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));

    std::vector<float> outputDx(reinterpret_cast<float *>(dxHost), reinterpret_cast<float *>(dxHost) + eleNum);

    CleanUp(dxHost, dyDevice, xDevice, dxDevice);
    CleanACL(stream, context, deviceId);

    if (!VerifyResults(golden, outputDx)) {
            return -1;
    }
    printf("Accuracy verification passed.\n");
    return 0;
}
